import pickle
import random

import trimesh
from autolab_core import RigidTransform
from copy import deepcopy
import numpy as np
import json
import os
#from .grasping.grasp import ParallelJawPtGrasp3D
# from grasping.grasp_sampler import GraspSampler

cur_dir = os.path.dirname(__file__)
object_dict_path = os.path.join(cur_dir, '../grasp_suction_dataset/object_models/objects_dict.json')
object_grasp_path = os.path.join(cur_dir, '../grasp_suction_dataset/grasp_10000')
finger_path = os.path.join(cur_dir, '../grippers/yumi/coarse/wider_finger.stl')


def change_ray_color(ray, color):
    colors = np.ones((len(ray.entities), 3))
    colors_1 = (colors * color).astype(np.uint8)  # approach
    ray.colors = colors_1


# np.random.seed(0)
# random.seed(0)

T = trimesh.transformations.translation_matrix([-0.002, -0.0065, -0.0480])
R_l = trimesh.transformations.euler_matrix(np.pi/2, np.pi/2, 0, 'rxyz')
R_r = trimesh.transformations.euler_matrix(np.pi/2, np.pi/2, np.pi, 'rxyz')

with open(object_dict_path, 'r') as f:
    dataset_dict = json.load(fp=f)
    for key in dataset_dict.keys():
        # if not key.startswith('3dnet_102'):
        #    continue
        path = dataset_dict[key]['path']
        mesh = trimesh.load_mesh(path)
        trimesh.smoothing.filter_humphrey(mesh)
        files = os.listdir('{}/{}'.format(object_grasp_path, key))
        for file_name in files:
            if os.path.splitext(file_name)[0].startswith('final'):
                with open('{}/{}/{}'.format(object_grasp_path, key, file_name), 'rb') as f:
                    grasps = pickle.load(f)
                    count = 0
                    print(key)
                    samples = random.sample(grasps, 10)
                    scene = trimesh.Scene([mesh])
                    for grasp in samples:
                        # print('grasp ori width is ',grasp.jaw_width_)
                        print('grasp pad width is ', grasp.real_jaw_width_)
                        print('grasp new width is ', grasp.grasp_width_new + grasp.real_jaw_width_ - grasp.jaw_width_)
                        print('extend pad is ', grasp.real_jaw_width_ - grasp.jaw_width_)
                        print('offset_distance',grasp.offset_distance)
                        # exit()
                        # print(grasp.center,grasp.center_offset,grasp.offset_distance,grasp.real_jaw_width_,grasp.grasp_width_new)
                        # exit()
                        if True:
                    # for i in range(int(len(grasps)/1)):
                    # for i in range(2):
                    #     grasp_ = grasps[i]
                    #     for grasp in grasps:
                            # print(grasp)
                            count+=1
                            distance = np.linalg.norm(grasp.center - np.array([0.00, 0, 0.01]))
                            #if grasp.grasp_score < 0.60:
                            angle = np.arccos(grasp.axis_.dot(np.array([1, 0, 0]))) / np.pi * 180
                            #if distance < 0.035 and  (60<angle < 120):
                            grasp_score = grasp.grasp_score.sum()
                            if grasp_score > 0:
                            # if grasp.real_jaw_width_  <  0.01:
                            #     print('jaw_width is : ', grasp.real_jaw_width_)
                                finger_left = trimesh.load(finger_path)
                                finger_left.visual.face_colors = [255,0,165,200]

                                finger_right = deepcopy(finger_left)

                                finger_left_offset = deepcopy(finger_left)
                                finger_left_offset.visual.face_colors = [255, 0, 255, 200]

                                finger_right_offset = deepcopy(finger_left_offset)

    #    axis, center, approach_angle = grasp axis, grasp.center, grasp.approach_angle
    #                            T_ = trimesh.transformations.translation_matrix([-grasp.real_jaw_width_/2, -0., -0.])
                                T_ = trimesh.transformations.translation_matrix([-0.002, -0.0065, -0.0480])
                                # print(grasp.approach_distance)
                                # finger_left.apply_transform(T)
                                finger_left.apply_transform(T_)
    #                             finger_right.apply_transform(T)
                                finger_right.apply_transform(T_)
                                finger_left.apply_transform(R_l)
                                finger_right.apply_transform(R_r)

                                finger_left_offset.apply_transform(T_)
                                finger_right_offset.apply_transform(T_)
                                finger_left_offset.apply_transform(R_l)
                                finger_right_offset.apply_transform(R_r)

                                approach_method = True
                                if approach_method:
                                    T_leftfinger = RigidTransform(grasp.approached_full_axis, grasp.virtual_contactpoints[0],
                                                                  from_frame='left_finger', to_frame='obj')
                                    T_rightfinger = RigidTransform(grasp.approached_full_axis, grasp.virtual_contactpoints[1],
                                                                   from_frame='right_finger', to_frame='obj')
                                    T_leftfinger_offset = RigidTransform(grasp.approached_full_axis, grasp.virtual_contactpoints_offset[0],
                                                                  from_frame='T_leftfinger_offset', to_frame='obj')
                                    T_rightfinger_offset = RigidTransform(grasp.approached_full_axis, grasp.virtual_contactpoints_offset[1],
                                                                   from_frame='T_rightfinger_offset', to_frame='obj')
                                else:
                                    T_leftfinger = RigidTransform(grasp.rotated_full_axis, grasp.virtual_contactpoints[0],
                                                                  from_frame='left_finger', to_frame='obj')
                                    T_rightfinger = RigidTransform(grasp.rotated_full_axis, grasp.virtual_contactpoints[1],
                                                                   from_frame='right_finger', to_frame='obj')

                                # ray_origins = np.array([[0, 0, 0]])
                                # ray_directions = np.array([[0, -1, 0]])
                                # ray_visualize = trimesh.load_path(np.hstack((ray_origins, ray_origins+ray_directions/10)).reshape(-1, 2, 3))
                                # scene_ = trimesh.Scene([ray_visualize, finger_left, finger_right])
                                # # scene_.show()
                                pc_vir_conpts_offset0 = trimesh.PointCloud(grasp.virtual_contactpoints_offset[0].reshape(1,3),colors=[255,165,0])
                                pc_vir_conpts_offset1 = trimesh.PointCloud(grasp.virtual_contactpoints_offset[1].reshape(1,3),
                                                              colors=[255, 165, 0])
                                pc_sampled_point = trimesh.PointCloud(grasp.center.reshape(1,3),colors=[0,255,0])
                                ray_normal = trimesh.load_path(
                                    np.hstack((grasp.center, grasp.center + grasp.grasp_approach / 10)).reshape(-1, 2, 3))
                                ray_axis = trimesh.load_path(np.hstack((
                                    grasp.center.reshape(1, 3),
                                    grasp.center.reshape(1, 3) + grasp.axis / 20)).reshape(-1, 2, 3))
                                ray_axis_ = trimesh.load_path(np.hstack((
                                    grasp.center.reshape(1, 3),
                                    grasp.center.reshape(1, 3) - grasp.axis / 20)).reshape(-1, 2, 3))
                                change_ray_color(ray_normal,[0,0,255])
                                change_ray_color(ray_axis_, [255, 0, 0])
                                change_ray_color(ray_axis, [255, 0, 0])
                                finger_left.apply_transform(T_leftfinger.matrix)
                                finger_right.apply_transform(T_rightfinger.matrix)
                                finger_left_offset.apply_transform(T_leftfinger_offset.matrix)
                                finger_right_offset.apply_transform(T_rightfinger_offset.matrix)
                                scene.add_geometry(finger_left)
                                scene.add_geometry(finger_right)
                                scene.add_geometry(pc_sampled_point)
                                scene.add_geometry(finger_left_offset)
                                scene.add_geometry(finger_right_offset)
                                scene.add_geometry([pc_vir_conpts_offset0,pc_vir_conpts_offset1])
                                # scene.add_geometry(ray_normal)
                                # scene.add_geometry(ray_axis_)
                                # scene.add_geometry(ray_axis)

                    print(count)
                    scene.show()

